from collections import defaultdict
from typing import Dict

from fastapi import WebSocket
from fastapi.websockets import WebSocketState

from app.utils.decorators.singleton import singleton


@singleton
class WebSocketConnectionManager:
    def __init__(self):
        self.active_connections: Dict[str, WebSocket] = {}
        self.user_groups: Dict[str, set] = defaultdict(set)

    async def connect(self, websocket: WebSocket, username: str):
        """连接WebSocket并添加到活跃连接"""
        if username in self.active_connections:
            # 如果用户已连接，先断开旧连接
            old_websocket = self.active_connections[username]
            try:
                if old_websocket.client_state == WebSocketState.CONNECTED:
                    await old_websocket.close(code=1000)
            except:
                pass
            finally:
                if username in self.active_connections:
                    del self.active_connections[username]

        await websocket.accept()
        self.active_connections[username] = websocket

    async def disconnect(self, username: str):
        if username in self.active_connections:
            websocket = self.active_connections[username]
            try:
                if websocket.client_state == WebSocketState.CONNECTED:
                    await websocket.close(code=1000)
            except:
                pass
            finally:
                if username in self.active_connections:
                    del self.active_connections[username]

        # 清理用户组
        for group in self.user_groups.values():
            group.discard(username)

    def is_connected(self, username: str) -> bool:
        websocket = self.active_connections.get(username)
        return websocket is not None and websocket.client_state == WebSocketState.CONNECTED

    def disconnect_by_websocket(self, websocket: WebSocket):
        for username, ws in list(self.active_connections.items()):
            if ws == websocket:
                del self.active_connections[username]
                break

    async def send_text(self, message: str, websocket: WebSocket):
        await websocket.send_text(message)

    async def send_json(self, data: dict, websocket: WebSocket):
        try:
            if websocket.client_state == WebSocketState.CONNECTED:
                await websocket.send_json(data)
        except Exception:
            # 记录错误或从活跃连接中移除
            self.disconnect_by_websocket(websocket)

    async def send_bytes(self, data: bytes, websocket: WebSocket):
        await websocket.send_bytes(data)

    async def send_private_json(self, data: dict, recipient: str) -> bool:
        if recipient in self.active_connections:
            await self.send_json(data, self.active_connections[recipient])
            return True
        return False

    async def broadcast_json(self, data: dict):
        disconnected = []
        for username, websocket in self.active_connections.items():
            try:
                if websocket.client_state == WebSocketState.CONNECTED:
                    await self.send_json(data, websocket)
                else:
                    disconnected.append(username)
            except:
                disconnected.append(username)

        # 清理已断开的连接
        for username in disconnected:
            await self.disconnect(username)

    async def broadcast_to_group(self, group_name: str, data: dict):
        if group_name in self.user_groups:
            for username in self.user_groups[group_name]:
                if username in self.active_connections:
                    await self.send_json(data, self.active_connections[username])

    async def add_to_group(self, username: str, group_name: str):
        self.user_groups[group_name].add(username)

    async def remove_from_group(self, username: str, group_name: str):
        self.user_groups[group_name].discard(username)